/*
 * Copyright (C) 2008 by The Regents of the University of California
 * Redistribution of this file is permitted under the terms of the GNU
 * Public License (GPL).
 *
 * @author Junghoo "John" Cho <cho AT cs.ucla.edu>
 * @date 3/24/2008
 */
 
#include "BTreeIndex.h"
#include "BTreeNode.h"

using namespace std;

/*
 * BTreeIndex constructor
 */
BTreeIndex::BTreeIndex()
{
    rootPid = -1;
}

/*
 * Open the index file in read or write mode.
 * Under 'w' mode, the index file should be created if it does not exist.
 * @param indexname[IN] the name of the index file
 * @param mode[IN] 'r' for read, 'w' for write
 * @return error code. 0 if no error
 */
RC BTreeIndex::open(const string& indexname, char mode)
{
	int rc;
	
	if (rc = pf.open(indexname, mode) < 0)		// open pf
		return rc;
	
	char buffer[PageFile::PAGE_SIZE];
	memset(&buffer, 0, PageFile::PAGE_SIZE);						// clear buffer
	
	if (rc = pf.read(0, buffer) < 0)								// read pf into buffer
		return rc; 
	
	memcpy(&rootPid, &buffer, sizeof(PageId));						// set rootPid
	memcpy(&treeHeight, &buffer[sizeof(PageId)], sizeof(int));		// set treeHeight
	
	if (rc = pf.write(0, buffer) < 0)
		return rc;
	
	return 0;
}

/*
 * Close the index file.
 * @return error code. 0 if no error
 */
RC BTreeIndex::close()
{
	int rc;
	char buffer[PageFile::PAGE_SIZE];
	memset(&buffer, 0, PageFile::PAGE_SIZE);						// clear buffer
	
	memcpy(&buffer[0], &rootPid, sizeof(PageId));					
	memcpy(&buffer[sizeof(PageId)], &treeHeight, sizeof(int));		
	
	if (rc = pf.write(0, buffer) < 0)
		return rc;
	return pf.close();
}

RC BTreeIndex::insertHelper(int key, const RecordId& rid, PageId& pid, int height)
{
	int rc;
	
	// if we are at a leaf node
	if (height == treeHeight)
	{
		BTLeafNode leaf;
		if (rc = leaf.read(pid, pf) < 0)
			return rc;
		rc = leaf.insert(key, rid);
		if (rc == RC_NODE_FULL)		// if leaf overflow
		{
			// create new sibling leaf node
			BTLeafNode sibling;
			int siblingKey;
			PageId siblingPid = pf.endPid();
			if (rc = sibling.read(siblingPid, pf) < 0)
				return rc;
				
			// insert and split nodes
			if (rc = leaf.insertAndSplit(key, rid, sibling, siblingKey) < 0)
				return rc;
				
			// set next node pointers
			sibling.setNextNodePtr(leaf.getNextNodePtr());
			leaf.setNextNodePtr(siblingPid);
			
			// write to PageFile
			if (rc = leaf.write(pid, pf) < 0)
				return rc;
			if (rc = sibling.write(siblingPid, pf) < 0)
				return rc;
				
			pid = siblingPid;
			return siblingKey;
		}
		else
		{
			// write to PageFile
			if (rc = leaf.write(pid, pf) < 0)
				return rc;
			return 0;
		}
	}
	
	// if we are not at a leaf node
	else
	{
		BTNonLeafNode nleaf;
		PageId childPtrPid;
		if (rc = nleaf.read(pid, pf) < 0)
			return rc;
		
		// find proper child node to insert into
		if (rc = nleaf.locateChildPtr(key, childPtrPid) < 0)
			return rc;
			
		// recursively insert into child
		int siblingKey = insertHelper(key, rid, childPtrPid, height + 1);
		
		// if there was overflow
		if (siblingKey > 0)		
		{
			rc = nleaf.insert(siblingKey, childPtrPid);			// insert into non-leaf
			
			// if non-leaf overflow
			if (rc == RC_NODE_FULL)
			{
				// create new sibling leaf node
				BTNonLeafNode sibling;
				int midKey;
				PageId siblingPid = pf.endPid();
				if (rc = sibling.read(siblingPid, pf) < 0)
					return rc;
					
				// insert and split nodes
				if (rc = nleaf.insertAndSplit(key, childPtrPid, sibling, midKey) < 0)
					return rc;
				
				// write to PageFile
				if (rc = nleaf.write(pid, pf) < 0)
					return rc;
				if (rc = sibling.write(siblingPid, pf) < 0)
					return rc;
				
				// if this is the root node
				if (pid == rootPid)
				{
					// create new root
					BTNonLeafNode root;
					PageId oldRootPid = rootPid;						// save old root pid
					rootPid = pf.endPid();
					if (rc = root.read(rootPid, pf) < 0)
						return rc;
					root.initializeRoot(oldRootPid, key, siblingPid);		// initialize new root
					if (rc = root.write(rootPid, pf) < 0)
						return rc;
					treeHeight++;
				}
				
				// not the root node
				else
				{
					pid = siblingPid;
					return siblingKey;
				}
			}
			
			// no non-leaf overflow
			else
			{
				// write to PageFile
				if (rc = nleaf.write(pid, pf) < 0)
					return rc;
				return 0;
			}
		}
		
		// no overflow
		else
			return 0;
	}
}

/*
 * Insert (key, RecordId) pair to the index.
 * @param key[IN] the key for the value inserted into the index
 * @param rid[IN] the RecordId for the record being inserted into the index
 * @return error code. 0 if no error
 */
RC BTreeIndex::insert(int key, const RecordId& rid)
{
	int rc;
	
	// if there is only a root node
	if (treeHeight == 0)
	{
		// create a new root
		BTNonLeafNode root;
		PageId newRootPid = pf.endPid();
		if (rc = root.read(newRootPid, pf) < 0)
			return rc;
			
		// create leaf nodes
		BTLeafNode leftLeaf;
		BTLeafNode rightLeaf;
		PageId leftLeafPid = pf.endPid() + 1;
		PageId rightLeafPid = pf.endPid() + 2;
		if (rc = leftLeaf.read(leftLeafPid, pf) < 0)
			return rc;
		if (rc = rightLeaf.read(rightLeafPid, pf) < 0)
			return rc;
			
		leftLeaf.setNextNodePtr(rightLeafPid);					// connect leaf nodes
		rightLeaf.insert(key, rid);								// insert into right leaf
		root.initializeRoot(leftLeafPid, key, rightLeafPid);	// initialize root
		
		// write back to PageFile
		if (rc = root.write(newRootPid, pf) < 0)
			return rc;
		if (rc = leftLeaf.write(leftLeafPid, pf) < 0)
			return rc;
		if (rc = rightLeaf.write(rightLeafPid, pf) < 0)
			return rc;
		rootPid = newRootPid;
		treeHeight++;
		return 0;
	}

    return insertHelper(key, rid, rootPid, 0);		// call helper function
}

/*
 * Find the leaf-node index entry whose key value is larger than or 
 * equal to searchKey, and output the location of the entry in IndexCursor.
 * IndexCursor is a "pointer" to a B+tree leaf-node entry consisting of
 * the PageId of the node and the SlotID of the index entry.
 * Note that, for range queries, we need to scan the B+tree leaf nodes.
 * For example, if the query is "key > 1000", we should scan the leaf
 * nodes starting with the key value 1000. For this reason,
 * it is better to return the location of the leaf node entry 
 * for a given searchKey, instead of returning the RecordId
 * associated with the searchKey directly.
 * Once the location of the index entry is identified and returned 
 * from this function, you should call readForward() to retrieve the
 * actual (key, rid) pair from the index.
 * @param key[IN] the key to find.
 * @param cursor[OUT] the cursor pointing to the first index entry
 *                    with the key value.
 * @return error code. 0 if no error.
 */
RC BTreeIndex::locate(int searchKey, IndexCursor& cursor)
{
	int rc;
	BTNonLeafNode nleaf;
	PageId pid = rootPid;
	
	// find leaf node with searchKey
	for (int i = treeHeight; i > 0; i--)
	{
		if (rc = nleaf.read(pid, pf) < 0)
			return rc;
		if (rc = nleaf.locateChildPtr(searchKey, pid) < 0)
			return rc;
	}
	
	// retrieve eid from leaf node
	BTLeafNode leaf;
	int eid;
	if (rc = leaf.read(pid, pf) < 0)
		return rc;
	if (leaf.locate(searchKey, eid) < 0)
		return rc;
	
	// return pid and eid to cursor
	cursor.pid =  pid;
	cursor.eid = eid;
	
    return 0;
}

/*
 * Read the (key, rid) pair at the location specified by the index cursor,
 * and move foward the cursor to the next entry.
 * @param cursor[IN/OUT] the cursor pointing to an leaf-node index entry in the b+tree
 * @param key[OUT] the key stored at the index cursor location.
 * @param rid[OUT] the RecordId stored at the index cursor location.
 * @return error code. 0 if no error
 */
RC BTreeIndex::readForward(IndexCursor& cursor, int& key, RecordId& rid)
{
	int rc;
	BTLeafNode leaf;
	if (rc = leaf.read(cursor.pid, pf) < 0)
		return rc;
	if (rc = leaf.readEntry(cursor.eid, key, rid) < 0)
		return rc;
	
	int entry_size = sizeof(RecordId) + sizeof(int);
	if (key > 0 && cursor.eid < PageFile::PAGE_SIZE/entry_size - 1)
		cursor.eid++;
	
	// if reached the last entry
	else		
	{
		char buffer[PageFile::PAGE_SIZE];
		memset(&buffer, 0, PageFile::PAGE_SIZE);						// clear buffer
		
		if (rc = pf.read(0, buffer) < 0)								// read pf into buffer
			return rc; 
		
		memcpy(&cursor.pid, &buffer[PageFile::PAGE_SIZE-sizeof(PageId)], sizeof(PageId));
		if (cursor.pid <= 0)
			return RC_INVALID_PID;
		cursor.eid = 0;
	}
		
    return 0;
}
